import tensorflow as tf
from tensorflow.keras.optimizers import RMSprop
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.preprocessing.image import ImageDataGenerator

train_dir = r'data'
validation_dir = r'testData'

train_datagen = ImageDataGenerator(
    rescale = 1. / 255,  # 图像像素的归一化
    rotation_range = 40,  # 通过旋转图像的角度扩大训练集
    width_shift_range = 0.2,  # 通过水平偏移图像来扩大训练集
    height_shift_range = 0.2,  # 通过垂直偏移图像来扩大训练集
    shear_range = 0.2,  # 通过随机裁剪的角度来扩大训练集
    zoom_range = 0.2,  # 通过随机缩放来扩大训练集
    horizontal_flip = True,  # 通过随机将一半图像水平翻转来扩大训练集
    fill_mode = 'nearest')  # 部分操作过后，产生的像素缺失，通过填充最近的像素来补全

test_datagen = ImageDataGenerator(rescale = 1. / 255)

train_generator = train_datagen.flow_from_directory(
    train_dir,  # 加载训练集图像的存储地址
    target_size = (224, 224),  # 将图像尺寸变换为224X224
    batch_size = 21,  # 设置每批次加载的图像
    class_mode = 'categorical')  # 任务为多分类

validation_generator = test_datagen.flow_from_directory(
    validation_dir,  # 加载测试集图像的存储地址用作验证
    target_size = (224, 224),  # 将图像尺寸变换为224X224
    batch_size = 21,  # 设置每批次加载的图像
    class_mode = 'categorical')  # 任务为多分类

model = tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(32, (10, 10), activation = 'relu', input_shape = (224, 224, 3)),  # 图像尺寸为224X224，3个颜色通道
    tf.keras.layers.MaxPooling2D(2, 2),
    tf.keras.layers.Conv2D(64, (3, 3), activation = 'relu'),
    tf.keras.layers.MaxPooling2D(2, 2),
    tf.keras.layers.Conv2D(128, (3, 3), activation = 'relu'),
    tf.keras.layers.MaxPooling2D(2, 2),
    tf.keras.layers.Conv2D(128, (3, 3), activation = 'relu'),
    tf.keras.layers.MaxPooling2D(2, 2),
    tf.keras.layers.Conv2D(256, (3, 3), activation = 'relu'),
    tf.keras.layers.MaxPooling2D(2, 2),
    tf.keras.layers.Conv2D(256, (3, 3), activation = 'relu'),
    tf.keras.layers.MaxPooling2D(2, 2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(512, activation = 'relu'),
    tf.keras.layers.Dense(4, activation = 'softmax')
])
model.compile(loss = 'categorical_crossentropy',  # 多分类任务
              optimizer = Adam(),  # 学习速率
              metrics = ['accuracy'])  # 设置损失函数为accuracy

history = model.fit(
    train_generator,  # 加载训练集的完整预处理部分
    steps_per_epoch = 100,  # 每步训练加载100张图片
    epochs = 200,  # 模型训练10步
    validation_data = validation_generator,  # 加载测试集的完整预处理部分用于验证
    validation_steps = 25,  # 每步加载25张测试集图片用于验证
    verbose = 1)
